"""
PaddleOCR模型迁移脚本
将系统默认位置的PaddleOCR模型迁移到项目的models文件夹
"""

import os
import shutil
import sys
from pathlib import Path


def main():
    """
    主函数 - 迁移PaddleOCR模型
    """
    print("🚀 开始迁移PaddleOCR模型...")
    
    # 项目根目录
    project_root = Path(__file__).parent.parent
    target_dir = project_root / "models" / "paddleocr"
    
    # 系统默认的PaddleOCR模型位置
    user_home = Path.home()
    possible_sources = [
        user_home / ".paddleocr",
        user_home / ".paddlex" / "official_models",
        user_home / ".cache" / "paddle",
        user_home / "AppData" / "Local" / "paddle" if os.name == 'nt' else None,
    ]
    
    # 过滤掉None值
    possible_sources = [p for p in possible_sources if p is not None]
    
    print(f"📁 目标目录: {target_dir}")
    print(f"🔍 搜索模型文件...")
    
    # 创建目标目录
    target_dir.mkdir(parents=True, exist_ok=True)
    (target_dir / "official_models").mkdir(parents=True, exist_ok=True)
    
    models_found = False
    
    # 搜索并迁移模型
    for source_dir in possible_sources:
        if source_dir.exists():
            print(f"   检查: {source_dir}")
            
            # 查找模型文件
            model_files = []
            for pattern in ["*.pdiparams", "*.pdmodel", "*.yml", "*.json"]:
                model_files.extend(source_dir.rglob(pattern))
            
            if model_files:
                print(f"   ✅ 找到 {len(model_files)} 个模型文件")
                models_found = True
                
                # 迁移文件
                for model_file in model_files:
                    # 保持相对路径结构
                    relative_path = model_file.relative_to(source_dir)
                    target_path = target_dir / relative_path
                    
                    # 创建目标目录
                    target_path.parent.mkdir(parents=True, exist_ok=True)
                    
                    # 复制文件
                    if not target_path.exists():
                        try:
                            shutil.copy2(model_file, target_path)
                            print(f"   📋 复制: {relative_path}")
                        except Exception as e:
                            print(f"   ❌ 复制失败: {relative_path} - {e}")
                    else:
                        print(f"   ⏭️  跳过: {relative_path} (已存在)")
            else:
                print(f"   ❌ 未找到模型文件")
    
    if not models_found:
        print("⚠️  未找到任何PaddleOCR模型文件")
        print("💡 这可能是因为：")
        print("   1. 模型还未下载")
        print("   2. 模型存储在其他位置")
        print("   3. 使用了不同版本的PaddleOCR")
        print("\n🔧 建议解决方案：")
        print("   1. 运行程序让PaddleOCR自动下载模型")
        print("   2. 手动查找模型文件位置")
        print("   3. 重新安装PaddleOCR")
    else:
        print(f"✅ 模型迁移完成！")
        print(f"📍 模型现在位于: {target_dir}")
        
        # 显示迁移后的目录结构
        print("\n📂 迁移后的目录结构:")
        for root, dirs, files in os.walk(target_dir):
            level = root.replace(str(target_dir), '').count(os.sep)
            indent = ' ' * 2 * level
            print(f"{indent}{os.path.basename(root)}/")
            subindent = ' ' * 2 * (level + 1)
            for file in files:
                print(f"{subindent}{file}")
    
    print("\n🎯 下次运行程序时，PaddleOCR将使用项目目录中的模型")


if __name__ == "__main__":
    main() 